package redis

import (
	"crypto/sha1"
	"fmt"
	"github.com/gomodule/redigo/redis"
	"sieve_admin_server_userpass/lib/log"
	"strings"
	"time"
)

type Script struct {
	hash   string
	script string
}

func NewScript(script string) *Script {
	return &Script{script: script, hash: fmt.Sprintf("%x", sha1.Sum([]byte(script)))}
}

func (k *key) Eval(script *Script, args ...any) Result {
	conn := k.c.pool.Get()
	defer conn.Close()

	ss := []any{len(k.keys)}
	for _, key := range k.keys {
		ss = append(ss, key)
	}
	ss = append(ss, args...)

	reply, err := conn.Do("EVALSHA", append([]any{script.hash}, ss...)...)
	if e, ok := err.(redis.Error); ok && strings.HasPrefix(string(e), "NOSCRIPT ") {
		reply, err = conn.Do("EVAL", append([]any{script.script}, ss...)...)
	}
	return &cmd{reply: reply, err: err}
}

// Lock

type LockKey interface {
	Lock(timeout int) *Lock
}

func (c *Client) Lock(key ...string) LockKey {
	return c.key(key)
}

type Lock struct {
	key *key
	val string
}

func (k *key) Lock(timeout int) *Lock {
	val := time.Now().String()
	ok, err := k.Set(&SetArgs{
		Value: val,
		EX:    timeout,
		NX:    true,
	}).Bool()
	if err != nil || !ok {
		return nil
	}
	return &Lock{k, val}
}

var rsDelEq = NewScript(`if redis.call("get",KEYS[1])==ARGV[1] then return redis.call("del",KEYS[1]) else return 0 end`)

func (r *Lock) Unlock() bool {
	ok, err := r.key.Eval(rsDelEq, r.val).Bool()
	if err != nil {
		log.Err("Unlock:", err)
		return false
	}
	return ok
}

// RateLimit

type RateLimitKey interface {
	TokenBucket(cap int, genPeriod time.Duration, genCount int, take int) *TokenBucketState
}

func (c *Client) RateLimit(key ...string) RateLimitKey {
	return c.key(key)
}

var rsRateLimit = NewScript(`
local key = KEYS[1]
local now = tonumber(ARGV[1])
local cap = tonumber(ARGV[2])
local interval = tonumber(ARGV[3])
local take = tonumber(ARGV[4])
local ret = redis.call("hmget", key, "c", "t")
local c = tonumber(ret[1])
local t = tonumber(ret[2])
if t and c then
	if now > t + interval then
		local add = math.floor((now - t) / interval)
		c = c + add
		t = t + add * interval
	end
else
	t = now
	c = 1
end
if c > cap then c = cap end
if c < take then
	redis.call("hmset", key, "c", string.format("%d",c), "t", string.format("%d",t))
	return {0, c, t}
end
c = c - take
redis.call("hmset", key, "c", string.format("%d",c), "t", string.format("%d",t))
return {1, c, t}
`)

func (k *key) TokenBucket(cap int, genPeriod time.Duration, genCount int, take int) *TokenBucketState {
	cap = check_positive_int(cap)
	genCount = check_positive_int(genCount)
	take = check_positive_int(take)
	interval := check_positive_int64(genPeriod.Nanoseconds() / int64(genCount))
	state := &TokenBucketState{
		cap:      cap,
		period:   genPeriod,
		count:    genCount,
		interval: interval,
		take:     take,
		now:      time.Now().UnixNano(),
	}
	result, err := k.Eval(rsRateLimit, state.now, cap, interval, take).Int64s()
	if err != nil {
		state.err = err
		return state
	}
	if len(result) != 3 {
		state.err = fmt.Errorf("ratelimit return wrong format")
		return state
	}
	state.result = result[0]
	state.remain = result[1]
	state.genTime = result[2]
	return state
}

func check_positive_int(n int) int {
	if n < 1 {
		return 1
	}
	return n
}

func check_positive_int64(n int64) int64 {
	if n < 1 {
		return 1
	}
	return n
}

type TokenBucketState struct {
	cap      int           // 令牌桶容量
	period   time.Duration // 令牌生产周期
	count    int           // 令牌周期内生产数量
	interval int64         // 每个令牌生产周期
	take     int
	now      int64
	result   int64 // 获取令牌结果
	remain   int64 // 桶内剩余令牌
	genTime  int64 // 最近生产令牌时间
	err      error
}

func (s *TokenBucketState) Bool() (bool, error) {
	if s.err != nil {
		return false, s.err
	}
	return s.result != 0, nil
}

func (s *TokenBucketState) Remaining() (int, error) {
	if s.err != nil {
		return 0, s.err
	}
	return int(s.remain), nil
}

func (s *TokenBucketState) NextGenAfter() (time.Duration, error) {
	if s.err != nil {
		return 0, s.err
	}
	return time.Duration(s.genTime + s.interval - s.now), nil
}
